#include "stm32h7xx.h"
#include "Serial.h"

#define KALMAN_Q 0.01f
#define KALMAN_R 0.1f
#define FILTER_WINDOW 5

typedef struct {
    float TVOC;
    float CH2O;
    float CO2;
    float O2;
    float CO;
    float CH4;
    uint8_t status;
    uint32_t timestamp;
} AirQualityData;

typedef struct {
    float q;
    float r;
    float x;
    float p;
    float k;
} KalmanFilter;

typedef struct {
    float buffer[FILTER_WINDOW];
    uint8_t index;
    uint8_t count;
    float sum;
} MovingAverageFilter;

AirQualityData airData;
uint8_t Air_Data[32];

static KalmanFilter tvocFilter, ch2oFilter, co2Filter, o2Filter, coFilter, ch4Filter;
static MovingAverageFilter tvocMAFilter, ch2oMAFilter, co2MAFilter, o2MAFilter, coMAFilter, ch4MAFilter;

void KalmanFilter_Init(KalmanFilter* filter, float initialValue) {
    filter->q = KALMAN_Q;
    filter->r = KALMAN_R;
    filter->x = initialValue;
    filter->p = 1.0f;
}

float KalmanFilter_Update(KalmanFilter* filter, float measurement) {
    filter->p = filter->p + filter->q;
    filter->k = filter->p / (filter->p + filter->r);
    filter->x = filter->x + filter->k * (measurement - filter->x);
    filter->p = (1 - filter->k) * filter->p;
    return filter->x;
}

void MovingAverageFilter_Init(MovingAverageFilter* filter) {
    for (uint8_t i = 0; i < FILTER_WINDOW; i++) {
        filter->buffer[i] = 0.0f;
    }
    filter->index = 0;
    filter->count = 0;
    filter->sum = 0.0f;
}

float MovingAverageFilter_Update(MovingAverageFilter* filter, float value) {
    filter->sum -= filter->buffer[filter->index];
    filter->buffer[filter->index] = value;
    filter->sum += value;
    filter->index = (filter->index + 1) % FILTER_WINDOW;
    if (filter->count < FILTER_WINDOW) {
        filter->count++;
    }
    return filter->sum / filter->count;
}

void InitFilters() {
    KalmanFilter_Init(&tvocFilter, 0.0f);
    KalmanFilter_Init(&ch2oFilter, 0.0f);
    KalmanFilter_Init(&co2Filter, 400.0f);
    KalmanFilter_Init(&o2Filter, 20.9f);
    KalmanFilter_Init(&coFilter, 0.0f);
    KalmanFilter_Init(&ch4Filter, 0.0f);
    
    MovingAverageFilter_Init(&tvocMAFilter);
    MovingAverageFilter_Init(&ch2oMAFilter);
    MovingAverageFilter_Init(&co2MAFilter);
    MovingAverageFilter_Init(&o2MAFilter);
    MovingAverageFilter_Init(&coMAFilter);
    MovingAverageFilter_Init(&ch4MAFilter);
}

bool ValidateAirData() {
    if (Air_Data[0] != 0xFF || Air_Data[1] != 0x17) return false;
    if ((Air_Data[6] == 0xFF && Air_Data[7] == 0xFF) || 
        (Air_Data[8] == 0xFF && Air_Data[9] == 0xFF))
        return false;
    return true;
}

void Air_Process() {
    static bool filtersInitialized = false;
    if (!filtersInitialized) {
        InitFilters();
        filtersInitialized = true;
    }
    
    airData.timestamp = HAL_GetTick();
    
    if (!ValidateAirData()) {
        airData.status |= 0x01;
        return;
    }
    
    airData.status = 0x00;
    
    float rawTVOC = ((uint16_t)Air_Data[2] << 8) | Air_Data[3];
    float rawCH2O = ((uint16_t)Air_Data[4] << 8) | Air_Data[5];
    float rawCO2 = ((uint16_t)Air_Data[6] << 8) | Air_Data[7];
    float rawO2 = ((uint16_t)Air_Data[8] << 8) | Air_Data[9];
    float rawCO = ((uint16_t)Air_Data[10] << 8) | Air_Data[11];
    float rawCH4 = ((uint16_t)Air_Data[12] << 8) | Air_Data[13];
    
    rawTVOC *= 0.001f;
    rawCH2O *= 0.001f;
    rawO2 *= 0.01f;
    rawCO *= 0.1f;
    rawCH4 *= 0.01f;
    
    airData.TVOC = KalmanFilter_Update(&tvocFilter, rawTVOC);
    airData.CH2O = KalmanFilter_Update(&ch2oFilter, rawCH2O);
    airData.CO2 = KalmanFilter_Update(&co2Filter, rawCO2);
    airData.O2 = KalmanFilter_Update(&o2Filter, rawO2);
    airData.CO = KalmanFilter_Update(&coFilter, rawCO);
    airData.CH4 = KalmanFilter_Update(&ch4Filter, rawCH4);
    
    if (airData.CO2 > 5000) airData.status |= 0x02;
    if (airData.O2 < 19.5f || airData.O2 > 23.5f) airData.status |= 0x04;
}
